'''
统计量填充，包括，上一个有效值，均值
'''

import sys
sys.path.append(".")
sys.path.append('..')
from impute_method.impute_base import ImputeBase
from lib.data import MeanStdScalar


import pandas as pd
import numpy as np


class LastImpute(ImputeBase):
    def __init__(self, **kwargs):
        super(LastImpute, self).__init__(**kwargs)
        self.method = 'last'

    def impute(self):
        df = pd.DataFrame(self.masked_data[:, 0, :])
        df = df.where(df != 0, np.nan)
        df.interpolate(method='pad', inplace=True)  # 用上一个值填充，开头的nan不能被插值
        df.interpolate(method='backfill', inplace=True)  # 用后一个值填充开头的nan
        df.dropna(inplace=True)  # 导致了错位
        impute_data = df.values
        return self.devide_window(impute_data)


class MeanImpute(ImputeBase):
    def __init__(self, **kwargs):
        super(MeanImpute, self).__init__(**kwargs)
        self.method = 'mean'

    def impute(self):
        # 用原始数据计算均值
        mean = np.mean(self.ori_data[:, 0, :], 0)  # [n_samples, seq_len, feature_dim]，这有一丁点的误差，因为最后几个数据没有统计到，
        # scalar = MeanStdScalar(self.indicators)
        # mean, _ = scalar.get_statistics()
        impute_data = self.masked_data[:, 0, :]  # [n_samples, feature_dim]
        feature_dim = 4
        for i in range(feature_dim):
            # impute_data[:, i] =
            # impute_data[impute_data[:,i] == 0][i] = mean[i]
            for j in range(len(impute_data)):
                if impute_data[j][i] == 0:
                    impute_data[j][i] = mean[i]
        return self.devide_window(impute_data)


if __name__ == '__main__':
    kwargs = dict()
    kwargs['data_name'] = 'water'
    kwargs['data_type'] = 'masked'
    # kwargs['indicators'] = 'WATER_TEMPERATURE,AMMONIA,DISSOLVED_OXYGEN,TOTAL_NITROGEN'
    # kwargs['masked_indicator'] = 'TOTAL_NITROGEN'
    kwargs['indicators'] = 'WATER_TEMPERATURE,PH_VALUE,TOTAL_NITROGEN,DISSOLVED_OXYGEN'
    kwargs['masked_indicator'] = ''
    kwargs['seq_len'] = 24

    last_impute = LastImpute(**kwargs)
    last_impute.metirc()

    mean_impute = MeanImpute(**kwargs)
    mean_impute.metirc()

